package cc.magicjson.example.demo.thead.scheduler;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

public class ConcurrentTasksWithRejectionHandler {

    public static void main(String[] args) {
        int totalTasks = 20;
        int corePoolSize = 5;
        int maxPoolSize = 10;
        int queueCapacity = 2;

        // 创建一个CountDownLatch来确保所有任务同时开始
        CountDownLatch startLatch = new CountDownLatch(1);

        // 创建一个CountDownLatch来等待所有任务完成
        CountDownLatch endLatch = new CountDownLatch(totalTasks);

        // 创建一个AtomicInteger来记录被拒绝的任务数量
        AtomicInteger rejectedTaskCounter = new AtomicInteger(0);

        // 自定义拒绝策略
        RejectedExecutionHandler rejectionHandler = (r, executor) -> {
            int rejectedTaskNumber = rejectedTaskCounter.incrementAndGet();
            System.out.println("任务 " + rejectedTaskNumber + " 被拒绝");
            if(!executor.isShutdown()) {
                executor.getQueue().poll();
                executor.execute(r);
            }

        };

        // 创建一个具有核心5个线程，最大10个线程，任务队列大小为2，自定义拒绝策略的线程池
        ThreadPoolExecutor executor = new ThreadPoolExecutor(
            corePoolSize,        // 核心线程数
            maxPoolSize,         // 最大线程数
            60,                  // 空闲线程存活时间
            TimeUnit.SECONDS,    // 时间单位
            new ArrayBlockingQueue<>(queueCapacity), // 任务队列大小
            rejectionHandler     // 自定义拒绝策略
        );

        // 提交任务
        for (int i = 1; i <= totalTasks; i++) {
            final int taskId = i;
            executor.submit(() -> {
                try {
                    rejectedTaskCounter.incrementAndGet();
                    startLatch.await(); // 等待开始信号
                    System.out.println("任务 " + taskId + " 开始执行");
                    Thread.sleep((long) (Math.random() * 1000)); // 模拟任务执行时间
                    System.out.println("任务 " + taskId + " 执行完毕");
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    System.out.println("任务 " + taskId + " 被中断");
                } finally {
                    endLatch.countDown(); // 任务完成，计数器减1
                }
            });
        }

        System.out.println("所有任务准备就绪，开始执行...");
        startLatch.countDown(); // 启动所有任务

        try {
            endLatch.await(); // 等待所有任务完成
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        executor.shutdown();
		try {
			if(executor.awaitTermination(5, TimeUnit.SECONDS)) {
				executor.shutdownNow();
                System.out.println("所有任务执行完毕");
			}
		} catch (InterruptedException e) {
			Thread.currentThread().interrupt();
		}
	}
}
